package com.guandata.cdc.source.impl;

import com.guandata.cdc.helper.ParamsHelper;
import com.guandata.cdc.model.source.ConnectOptionsLoader;
import com.guandata.cdc.model.source.SourceConnectParam;
import com.guandata.cdc.model.DataTableType;
import com.guandata.cdc.source.GuanDataSourceFactory;
import com.guandata.cdc.model.db_type.postgres.PostgresSourceConfig;
import io.debezium.connector.postgresql.connection.PostgresConnection;
import io.debezium.jdbc.JdbcConfiguration;
import lombok.SneakyThrows;
import lombok.extern.slf4j.Slf4j;
import org.apache.flink.connector.jdbc.postgres.database.catalog.PostgresCatalog;
import org.apache.flink.table.api.Schema;
import org.apache.flink.table.api.Table;
import org.apache.flink.table.api.TableDescriptor;
import org.apache.flink.table.api.bridge.java.StreamTableEnvironment;
import org.apache.flink.table.catalog.ObjectPath;

import java.sql.ResultSet;
import java.sql.SQLException;

/**
 * Postgres使用Replica Slot进行数据复制,需要在PG配置文件中开启以下功能:
 * <pre>
 *   postgres
 *     - "-c"
 *     - "wal_level=logical"
 *     - "-c"
 *     - "max_wal_senders=xx"
 *     - "-c"
 *     - "max_replication_slots=xx"
 * </pre>
 * <pre>
 *     wal_level
 *     max_wal_senders
 *     max_replication_slots
 * </pre>
 * @see <a href="https://www.postgresql.org/docs/current/logical-replication-config.html">Postgres Logical Replication Configuration</a>
 */
@Slf4j
public class PostgresSQLSourceFactory extends GuanDataSourceFactory<PostgresSourceConfig> {
    private final static String QUERY_SLOT_EXIST = "SELECT 1 FROM pg_replication_slots WHERE slot_name = '%s';";
    private final static String CREATE_LOGIC_SLOT = "SELECT pg_create_logical_replication_slot('%s', '%s');";
    private final static String DROP_LOGIC_SLOT = "SELECT pg_drop_replication_slot('%s');";
    private final static String SHOW_WAL_LEVEL = "SHOW wal_level;";
    private final static String QUERY_SLOT_ACTIVE = "SELECT active FROM pg_replication_slots WHERE slot_name = '%s';";
    private final static String SET_REPLICA_FULL = "ALTER TABLE %s.%s REPLICA IDENTITY FULL";
    private final static String IS_OPEN_FULL = "SELECT relreplident FROM pg_class WHERE relname = '%s';";

    private static final String LOGICAL = "logical";
    private static final String REPLICA_FULL = "f";

    public PostgresSQLSourceFactory(PostgresSourceConfig config) {
        super(config);
        initializeReplication();
        super.schema = initializeSchema();
    }

    @SneakyThrows
    private Schema initializeSchema() {
        var jdbcCatalog = new PostgresCatalog(
            Thread.currentThread().getContextClassLoader(),
            config.getCatalogName("postgres_source"),
            config.sourceDatabase,
            config.sourceUsername,
            config.sourcePassword,
            config.sourceJdbcUrl
        );
        return jdbcCatalog
            .getTable(new ObjectPath(config.sourceDatabase, config.sourceTable))
            .getUnresolvedSchema();
    }

    @SneakyThrows
    private void initializeReplication() {
        try(var connection = new PostgresConnection(buildConfiguration(), "postgres sql init")) {
            // 检查wal_level设置
            String walLevel = connection.queryAndMap(SHOW_WAL_LEVEL, rs -> {
                if (rs.next()) {
                    return rs.getString(1);
                }
                return null;
            });

            log.info("wal level: {}", walLevel);

            if (walLevel == null || !walLevel.equals(LOGICAL)) {
                throw new SQLException("Postgres SQL wal_level must be set to 'logical' for CDC. Current value: " + walLevel);
            }

            // 检查slot是否存在
            boolean slotExists = connection.queryAndMap(String.format(QUERY_SLOT_EXIST, config.slotName), ResultSet::next);

            log.info("slot {} exists: {}", config.slotName, slotExists);

            if (slotExists) {
                // 检查slot是否活跃
                Boolean isActive = connection.queryAndMap(String.format(QUERY_SLOT_ACTIVE, config.slotName), rs -> {
                    if (rs.next()) {
                        return rs.getBoolean(1);
                    }
                    return false;
                });

                if (!isActive) {
                    // 如果slot存在但不活跃，先删除它
                    log.info("Found inactive replication slot {}, dropping it", config.slotName);
                    connection.execute(String.format(DROP_LOGIC_SLOT, config.slotName));
                    slotExists = false;
                } else {
                    log.info("Found active replication slot {}, reusing it", config.slotName);
                }
            }

            // 如果slot不存在，则创建slot
            if (!slotExists) {
                log.info("Creating new replication slot: {}", config.slotName);
                connection.execute(String.format(CREATE_LOGIC_SLOT, config.slotName, config.pluginName));
            }

            // set replica full, 否则会导致没有before值，导致无法复制
            try {
                var isReplicaFull = connection.queryAndMap(String.format(IS_OPEN_FULL, config.rawSourceTable), rs -> {
                    if(rs.next()) {
                        return REPLICA_FULL.equalsIgnoreCase(rs.getString(1));
                    }
                    return false;
                });
                if(!isReplicaFull) {
                    connection.execute(String.format(SET_REPLICA_FULL, config.sourceSchema, config.rawSourceTable));
                }
            } catch (SQLException e) {
                log.error("NOTICE: Set replica failed!!!", e);
            }
        }
    }

    private JdbcConfiguration buildConfiguration() {
        return JdbcConfiguration
                .create()
                .withHostname(config.sourceHostname)
                .withPort(config.sourcePort)
                .withDatabase(config.sourceDatabase)
                .withConnectionTimeoutMs(60000)
                .withUser(config.sourceUsername)
                .withPassword(config.sourcePassword)
                .build();
    }

    @Override
    public Table registerSourceTable(StreamTableEnvironment tableEnv) {
        var sourceFlinkTableName = ParamsHelper.getInnerTableName(config.rawSourceTable);

        log.info("{} table's schema: {}", sourceFlinkTableName, this.schema);

        var descriptorBuilder = TableDescriptor.forConnector(config.connector).schema(this.schema);
        ConnectOptionsLoader
            .<SourceConnectParam>builder(DataTableType.POSTGRES, descriptorBuilder)
            .add(SourceConnectParam.HOST_NAME, config.sourceHostname)
            .add(SourceConnectParam.PORT, String.valueOf(config.sourcePort))
            .add(SourceConnectParam.DATABASE_NAME, config.sourceDatabase)
            .add(SourceConnectParam.SCHEMA_NAME, config.sourceSchema)
            .add(SourceConnectParam.TABLE_NAME, config.rawSourceTable)
            .add(SourceConnectParam.USERNAME, config.sourceUsername)
            .add(SourceConnectParam.PASSWORD, config.sourcePassword)
            .add(SourceConnectParam.SLOT_NAME, config.slotName)
            .add(SourceConnectParam.DECODING_PLUGIN_NAME, config.pluginName)
            .add(SourceConnectParam.SCAN_INCREMENTAL_SNAPSHOT_ENABLED, String.valueOf(config.isIncrementalMode()))
            .add(SourceConnectParam.SCAN_STARTUP_MODE, config.startUpMode.getCode())
            .add(SourceConnectParam.SCAN_INCREMENTAL_SNAPSHOT_CHUNK_SIZE, String.valueOf(8096));
        tableEnv.createTable(sourceFlinkTableName, descriptorBuilder.build());
        return dropColumnsFrom(tableEnv.from(sourceFlinkTableName));
    }
}
